{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Reservoir constraints\n",
    "To illustrate the usage of the reservoir constraint, we look at an example for scheduling nurses in clinic."
   ],
   "id": "3e38387cdf7f1298"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-20T21:16:05.468188Z",
     "start_time": "2024-08-20T21:16:04.867305Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from ortools.sat.python import cp_model\n",
    "\n",
    "model = cp_model.CpModel()"
   ],
   "id": "457bc28276122b9b",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The clinic needs to ensure that there are always enough nurses available without over-staffing too much.\n",
    "For a 12-hour work day, we model the demands for nurses as integers for each hour of the day. We also have a list of nurses, each with an individual availability as well as a maximum shift length.\n"
   ],
   "id": "111433800be0ed4a"
  },
  {
   "cell_type": "code",
   "id": "d1837cef744b2c56",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-08-20T21:16:05.472659Z",
     "start_time": "2024-08-20T21:16:05.469706Z"
    }
   },
   "source": [
    "demand_change_at_t = [3, 0, 0, 0, 2, 0, 0, 0, -1, 0, -1, 0, -3]\n",
    "demand_change_times = list(range(len(demand_change_at_t)))\n",
    "\n",
    "# begin and end of the availability of each nurse\n",
    "nurse_availabilities = 2 * [\n",
    "    (0, 7),\n",
    "    (0, 4),\n",
    "    (0, 8),\n",
    "    (2, 9),\n",
    "    (1, 5),\n",
    "    (5, 12),\n",
    "    (7, 12),\n",
    "    (0, 12),\n",
    "    (4, 12),\n",
    "]\n",
    "\n",
    "max_shift_length = 5"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "We now initialize all relevant variables of the model. Each nurse is assigned a start and end time of their shift as well as a Boolean variable indicating if they are working at all. We also add some basic constraints to ensure that the shifts are valid.",
   "id": "b4dc3e7d2e533022"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-20T21:16:07.398779Z",
     "start_time": "2024-08-20T21:16:07.394272Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# boolean variable to indicate if a nurse is scheduled\n",
    "nurse_scheduled = [\n",
    "    model.new_bool_var(f\"nurse_{i}_scheduled\") for i in range(len(nurse_availabilities))\n",
    "]\n",
    "\n",
    "# model the begin and end of each shift\n",
    "shifts_begin = [\n",
    "    model.new_int_var(begin, end, f\"begin_nurse_{i}\")\n",
    "    for i, (begin, end) in enumerate(nurse_availabilities)\n",
    "]\n",
    "\n",
    "shifts_end = [\n",
    "    model.new_int_var(begin, end, f\"end_nurse_{i}\")\n",
    "    for i, (begin, end) in enumerate(nurse_availabilities)\n",
    "]\n",
    "\n",
    "for begin, end in zip(shifts_begin, shifts_end):\n",
    "    model.add(end >= begin)  # make sure the end is after the begin\n",
    "    model.add(end - begin <= max_shift_length)  # make sure, the shifts are not too long"
   ],
   "id": "727db34f365ff385",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Our reservoir level is the number of nurses scheduled at any time minus the demand for nurses up until that point. We can now add the reservoir constraint to ensure that we have enough nurses available at all times while not having too many nurses scheduled (i.e., the reservoir level is between 0 and 2). We have three types of changes in the reservoir:\n",
    "\n",
    "1. The demand for nurses changes at the beginning of each hour. For these we use fixed integer times and activate all\n",
    "   changes. Note that the demand changes are negated, as an increase in demand lowers the reservoir level.\n",
    "2. If a nurse begins a shift, we increase the reservoir level by 1. We use the `shifts_begin` variables as times and\n",
    "   change the reservoir level only if the nurse is scheduled.\n",
    "3. Once a nurse ends a shift, we decrease the reservoir level by 1. We use the `shifts_end` variables as times and\n",
    "   change the reservoir level only if the nurse is scheduled."
   ],
   "id": "8ad34369fee29117"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-20T21:16:10.511025Z",
     "start_time": "2024-08-20T21:16:10.502550Z"
    }
   },
   "cell_type": "code",
   "source": [
    "times = demand_change_times\n",
    "demands = [\n",
    "    -demand for demand in demand_change_at_t\n",
    "]  # an increase in demand lowers the reservoir\n",
    "actives = [1] * len(demand_change_times)\n",
    "\n",
    "times += list(shifts_begin)\n",
    "demands += [1] * len(shifts_begin)  # a nurse begins a shift\n",
    "actives += list(nurse_scheduled)\n",
    "\n",
    "times += list(shifts_end)\n",
    "demands += [-1] * len(shifts_end)  # a nurse ends a shift\n",
    "actives += list(nurse_scheduled)\n",
    "\n",
    "model.add_reservoir_constraint_with_active(\n",
    "    times=times,\n",
    "    level_changes=demands,\n",
    "    min_level=0,\n",
    "    max_level=2,\n",
    "    actives=actives,\n",
    ")"
   ],
   "id": "444ecfc2f2078c79",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<ortools.sat.python.cp_model.Constraint at 0x10717b490>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Finally, we can solve the model and print the results.",
   "id": "6c1cf272286e79f9"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-20T21:16:14.600390Z",
     "start_time": "2024-08-20T21:16:14.125694Z"
    }
   },
   "cell_type": "code",
   "source": [
    "solver = cp_model.CpSolver()\n",
    "status = solver.solve(model)\n",
    "\n",
    "assert status == cp_model.OPTIMAL\n",
    "\n",
    "for i in range(len(nurse_availabilities)):\n",
    "    if solver.Value(nurse_scheduled[i]):\n",
    "        print(\n",
    "            f\"Nurse {i} is scheduled from {solver.Value(shifts_begin[i])} to {solver.Value(shifts_end[i])}\"\n",
    "        )"
   ],
   "id": "initial_id",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nurse 0 is scheduled from 2 to 7\n",
      "Nurse 1 is scheduled from 0 to 2\n",
      "Nurse 2 is scheduled from 3 to 8\n",
      "Nurse 3 is scheduled from 4 to 9\n",
      "Nurse 5 is scheduled from 5 to 10\n",
      "Nurse 6 is scheduled from 11 to 11\n",
      "Nurse 7 is scheduled from 7 to 12\n",
      "Nurse 8 is scheduled from 8 to 12\n",
      "Nurse 9 is scheduled from 0 to 5\n",
      "Nurse 10 is scheduled from 0 to 4\n",
      "Nurse 13 is scheduled from 3 to 5\n",
      "Nurse 14 is scheduled from 5 to 7\n",
      "Nurse 15 is scheduled from 11 to 11\n",
      "Nurse 17 is scheduled from 7 to 12\n"
     ]
    }
   ],
   "execution_count": 5
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
